import json
import os
import importlib
import yaml

# 导入损失函数工具
from training.utils.loss_util import create_loss_function

# 导入工具模块
from training.utils.optimizer_util import create_optimizer

# 导入统计工具
from training.utils.stats_util import get_class_weights_from_stats


def load_config(config_path):
    """加载配置文件"""
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def load_run_info(script_dir):
    """加载之前运行的信息"""
    run_info_path = os.path.join(script_dir, 'run_info.json')
    if os.path.exists(run_info_path):
        with open(run_info_path, 'r') as f:
            return json.load(f)
    return None


def save_run_info(script_dir, run_id, run_name, config):
    """保存运行信息"""
    run_info_path = os.path.join(script_dir, 'run_info.json')
    run_info = {
        'run_id': run_id,
        'run_name': run_name,
        'config': config
    }
    with open(run_info_path, 'w') as f:
        json.dump(run_info, f, indent=2)


def get_model(config):
    """
    从配置文件中动态加载模型
    
    参数:
    - config: 配置字典，包含model部分或arch部分
    
    返回值:
    - model: 加载的模型实例
    
    异常:
    - ValueError: 当模型文件不存在或无法加载时抛出
    """
    # 支持从model或arch部分加载模型配置
    if 'model' in config:
        model_section = config['model']
    elif 'arch' in config:
        model_section = config['arch']
    else:
        raise ValueError("配置文件中缺少 'model' 或 'arch' 字段")
    
    if 'name' not in model_section and 'type' not in model_section:
        raise ValueError("配置文件中缺少 'model.name' 或 'arch.type' 字段")
    
    # 获取模型类型
    model_type = model_section.get('name', model_section.get('type'))
    # 获取模型参数
    model_args = model_section.get('params', model_section.get('args', {}))
    
    # 特殊处理领域适应模型，确保它能从正确的文件加载
    model_filename = f"{model_type}.py"
    model_file_path = os.path.join(os.path.dirname(__file__), '..', 'models', model_filename)
    
    if not os.path.exists(model_file_path):
        raise ValueError(f"模型文件不存在: {model_file_path}")
    
    try:
        # 动态导入模型模块
        module_name = f"training.models.{model_type}"
        model_module = importlib.import_module(module_name)
        
        # 获取模型类
        model_class = getattr(model_module, model_type)
        
        # 创建模型实例
        model = model_class(**model_args)
        
        print(f"成功加载模型: {model_type}")
        return model
    except Exception as e:
        raise ValueError(f"加载模型 {model_type} 失败: {str(e)}")


def get_optimizer(model, config):
    """
    从配置文件中动态加载优化器（使用optimizer_util模块）
    
    参数:
    - model: 模型实例
    - config: 配置字典，包含optimizer部分
    
    返回值:
    - optimizer: 加载的优化器实例
    
    异常:
    - ValueError: 当优化器不存在或无法加载时抛出
    """
    if 'optimizer' not in config:
        raise ValueError("配置文件中缺少 'optimizer' 字段")
    
    # 使用新的optimizer_util模块创建优化器
    try:
        optimizer = create_optimizer(model, config['optimizer'])
        return optimizer
    except Exception as e:
        raise ValueError(f"加载优化器失败: {str(e)}")


def get_class_weights_from_config(config):
    """
    直接从配置文件中的统计文件路径获取类别权重，不依赖于数据集对象
    
    参数:
    - config: 配置字典
    
    返回值:
    - class_weights: 计算得到的类别权重列表，如果无法获取则返回None
    """
    # 尝试从配置中获取stats_file参数
    stats_file = None
    if 'data_loader' in config and config['data_loader'].get('stats_file'):
        stats_file = config['data_loader'].get('stats_file')
        print(f"从配置中获取统计文件路径: {stats_file}")
    
    # 使用工具函数直接从统计文件加载类别权重
    if stats_file:
        return get_class_weights_from_stats(stats_file)
    else:
        print("配置中未找到统计文件路径")
        return None

def get_loss_function(config, dataset=None, train_files=None):
    """
    从配置文件中动态加载损失函数，优先从统计文件获取权重
    
    参数:
    - config: 配置字典，包含loss部分
    - dataset: 可选，训练数据集对象（现在不再使用，保持向后兼容）
    - train_files: 可选，训练文件列表，用于从统计文件中筛选训练集数据计算权重
    返回值:
    - criterion: 加载的损失函数实例
    
    异常:
    - ValueError: 当损失函数不存在或无法加载时抛出
    """
    if 'loss' not in config:
        raise ValueError("配置文件中缺少 'loss' 字段")
    
    loss_type = config['loss']
    loss_args = config.get('loss_args', {})
    
    # 尝试从配置中获取stats_file参数
    stats_file = None
    if 'data_loader' in config and config['data_loader'].get('stats_file'):
        stats_file = config['data_loader'].get('stats_file')
        print(f"从配置中获取统计文件路径: {stats_file}")
    
    # 优化：当需要自动计算权重且有统计文件时，尝试先从统计文件获取权重
    if (loss_type == 'WeightedCrossEntropyLoss' and loss_args.get('weight') == 'auto') or \
       (loss_type == 'FocalLoss' and loss_args.get('alpha') == 'auto'):
        # 即使dataset为None，也尝试从统计文件获取权重
        if stats_file:
            print("提供了统计文件，从统计文件获取权重...")
            # 如果提供了train_files参数，则只使用训练集文件计算权重
            class_weights = get_class_weights_from_stats(stats_file, train_files)
            if class_weights is not None:
                # 设置权重参数
                if loss_type == 'WeightedCrossEntropyLoss':
                    loss_args['weight'] = class_weights
                elif loss_type == 'FocalLoss':
                    loss_args['alpha'] = class_weights
                print(f"成功从统计文件获取权重并设置到损失函数参数中")
    
    # 使用loss_util中的create_loss_function创建损失函数，传递stats_file参数
    # 不再依赖dataset参数，只从统计文件获取权重
    return create_loss_function(loss_type, loss_args, None, stats_file)


from training.utils.scheduler_util import create_scheduler

def get_lr_scheduler(optimizer, config):
    """
    从配置文件中动态加载学习率调度器（使用scheduler_util模块）
    
    参数:
    - optimizer: 优化器实例
    - config: 配置字典，包含lr_scheduler部分
    
    返回值:
    - scheduler_manager: 调度器管理器实例
    
    异常:
    - ValueError: 当学习率调度器不存在或无法加载时抛出
    """
    # 使用新的scheduler_util模块创建调度器和管理器
    try:
        scheduler_manager, _ = create_scheduler(optimizer, config)
        return scheduler_manager
    except Exception as e:
        raise ValueError(f"加载学习率调度器失败: {str(e)}")
